import threading
from MyTT import EMA
import numpy as np
import pandas as pd
from tqsdk import TqApi
from db_manager import connect_db, get_monitored_products  # 新增导入

class TqStrategy01:
    def __init__(self, api: TqApi, lots=1):
        self.api = api
        self.lots = lots
        self.running = False
        self.threads = []
        self.conn = connect_db()  # 使用标准数据库连接
        
        # 获取监控中的品种（使用项目标准方法）
        self.products = get_monitored_products(self.conn)  # 返回 (id, name, exchange_code,...)
    
    def start(self):
        """启动策略"""
        self.running = True
        # 为每个品种创建线程（过滤掉无交易所代码的品种）
        for product in self.products:
            if product[2]:  # exchange_code在元组索引2的位置
                thread = threading.Thread(target=self._process_product, args=(product,))
                self.threads.append(thread)
                thread.start()
    
    def _process_product(self, product):
        """单个品种处理线程"""
        symbol = product[2]  # 使用exchange_code作为合约代码
        klines = self.api.get_kline_serial(symbol, 86400, 100)
        
        while self.running:
            self.api.wait_update(deadline=0.05)
            if not klines.closed:
                # 计算技术指标
                q = (3*klines.close + klines.low + klines.open + klines.high) / 6
                ma_values = self._calculate_ma(q)
                short_ma = ma_values['ma']
                ema_7 = ma_values['ema']
                
                # 获取当前持仓
                position = self.api.get_position(symbol)
                
                # 修改后的交易逻辑
                if short_ma[-2] > ema_7[-2]:  # 做多信号
                    if position.pos_short > 0:
                        # 平全部空仓
                        self.api.insert_order(symbol, direction="BUY", offset="CLOSE", 
                                           volume=position.pos_short)
                    # 开多仓
                    self.api.insert_order(symbol, direction="BUY", offset="OPEN", 
                                       volume=self.lots)
                    
                elif short_ma[-2] < ema_7[-2]:  # 做空信号
                    if position.pos_long > 0:
                        # 平全部多仓
                        self.api.insert_order(symbol, direction="SELL", offset="CLOSE", 
                                           volume=position.pos_long)
                    # 开空仓
                    self.api.insert_order(symbol, direction="SELL", offset="OPEN", 
                                       volume=self.lots)
    
    def _calculate_ma(self, q):
        """自定义权重计算+MyTT优化"""
      
        # 保持原有的权重计算方式
        weights = np.arange(26, 0, -1)
        ma = []
        q_array = np.array(q)
        for i in range(len(q_array)):
            if i < 26:
                ma.append(np.nan)
                continue
            window = q_array[i-26:i+1]
            ma.append(np.dot(window, weights) / 351)
        
        # 使用MyTT的EMA计算
        ema_7 = EMA(pd.Series(ma), 7)
        
        return {'ma': ma, 'ema': ema_7.tolist()}
    
    def stop(self):
        """停止策略"""
        self.running = False
        for thread in self.threads:
            thread.join()
        # 关闭数据库连接
        if self.conn:
            self.conn.close()
